import abc
from ccxt.base.types import OrderSide
import ccxt.pro as ccxt
from collections import deque
from common.clickhouse_query import get_factors, get_klines
from common.variables import TimeFrameType
from copy import deepcopy
from trade.data.add_factors import add_batch_factors
from trade.data.add_kline_data import add_batch_klines
from typing import Dict, Union


RecordsType = dict[tuple[str, str], deque[float]]


class BrokerBase(abc.ABC):
    exchange: ccxt.Exchange
    timeframe: str
    cache_dir: str
    symbols: list
    positions: dict
    balances: dict
    markets: dict
    leverages: dict
    slip_rate: float
    data_records: RecordsType  # (symbol, fname): [1.2, 2.1, 3.4] etc.

    def reset(self):
        pass

    @abc.abstractmethod
    def prepare_sliding(
        self,
        start: str,
        end: str,
        need_symbols: list[str],
        need_fnames: list[str],
        timeframe: TimeFrameType,
    ):
        pass

    def get_data(self, name, col, idx: Union[slice, int]):
        if isinstance(idx, int):
            return self.data_records[(name, col)][idx]
        elif isinstance(idx, slice):
            t = self.data_records[(name, col)]
            s = idx.start
            end = idx.stop if idx.stop else 0
            p = idx.step if idx.step else 1
            results = []
            while s < end:
                results.append(t[s])
                s += p
            return results

    def init_get_symbol_and_factors(
        self,
        start: str,
        end: str,
        need_symbols: list[str],
        need_bfnames: list[str],
        timeframe: TimeFrameType,
    ):
        data_records = self.get_symbol_and_factors(
            start, end, need_symbols, need_bfnames, timeframe
        )
        self.data_records = data_records

    def get_symbol_and_factors(
        self,
        start: str,
        end: str,
        need_symbols: list[str],
        need_bfnames: list[str],
        timeframe: TimeFrameType,
    ):
        data_records = {}
        if not need_symbols:
            return data_records

        data_records: RecordsType

        add_batch_klines(need_symbols, start, end, timeframe, max_part=1000, n_jobs=4)

        klines = get_klines(start, end, need_symbols, timeframe)
        for symbol, df in klines.items():
            symbol: str
            if ("default", "time") not in data_records:
                data_records[("default", "time")] = deque(
                    (df["time"].astype("int64") / 1000000).astype("int64").tolist()
                )
            for col in ["open", "high", "low", "close", "volume"]:
                data_records[(symbol, col)] = deque(df[col].tolist())

        if need_bfnames:
            add_batch_factors(
                need_bfnames,
                need_symbols,
                start,
                end,
                timeframe,
                max_part=1000,
                n_jobs=4,
            )
            factors = get_factors(start, end, need_symbols, need_bfnames, timeframe)
            for symbol, df in factors.items():
                symbol_fname: tuple[str, str]
                for col in df.columns:
                    data_records[(symbol, col)] = deque(df[col], maxlen=len(df[col]))

        sizes = [len(x) for x in data_records.values()]
        if not sizes:
            return data_records

        if not all([x == sizes[0] for x in sizes]):
            raise Exception(
                f"初始化数据缺失：{start} {end} {need_symbols} {need_bfnames}"
            )

        return data_records

    @abc.abstractmethod
    def slide(self):
        pass

    def get_positions(self):
        return deepcopy(self.positions)

    def get_position_size(self, symbol: str):
        if symbol not in self.positions:
            return 0
        elif self.positions[symbol]["side"] == "long":
            return self.positions[symbol]["contracts"]
        else:
            return -self.positions[symbol]["contracts"]

    def get_balance_size(self, balance: str):
        if balance not in self.balances:
            return 0
        elif self.balances[balance]["total"] >= 0:
            return self.balances[balance]["total"]
        else:
            return -self.balances[balance]["debt"]

    def get_balances(self):
        return deepcopy(self.balances)

    def _get_cur_price(self, symbol: str):
        return self.get_data(symbol, "close", -1)

    def set_leverage(self, leverages: Dict[str, float]):
        self.leverages.update(leverages)

    def get_leverage(self, symbol: str):
        if symbol in self.leverages:
            return self.leverages[symbol]
        elif symbol in self.positions:
            return self.positions[symbol]["leverage"]
        else:
            return 1

    def get_balance_leverage(self, balance: str):
        return 10

    def get_collateral(self, balance: str):
        return 1

    def get_balance_margin(self, balance: str, debt: float = None, quote: str = "USDT"):
        if balance not in self.balances or self.balances["debt"] == 0:
            return 0
        if self.get_balance_leverage(balance) <= 1:
            raise Exception(
                f"{balance} leverage less than 1, but you try to calculate debt"
            )
        if not debt:
            debt = self.balances["debt"]
        return (
            debt
            * self._get_convert_ratio(balance, quote)
            / (self.get_balance_leverage(balance) - 1)
        )

    def _get_convert_ratio(self, base: str, quote: str) -> float:
        if quote == "USDT":
            if base == "USDT":
                return 1
            symbol = f"{base}/{quote}"
            price = self.get_data(symbol, "close", -1)
            return price
        else:
            price1 = self._get_convert_ratio(quote, "USDT")
            price2 = self._get_convert_ratio(base, "USDT")
            return price1 / price2

    def get_balance_value(self, balance: str):
        return (
            (self.balances[balance]["total"] - self.balances[balance]["debt"])
            * self.get_collateral(balance)
            * self._get_convert_ratio(balance, "USDT")
        )

    def get_all_balance_value(self):
        balance_value = 0
        for balance in self.balances:
            balance_value += self.get_balance_value(balance)
        return balance_value

    def get_symbol_unrealized(self, symbol: str):
        profit = (
            (self.get_data(symbol, "close", -1) - self.positions[symbol]["entryPrice"])
            * self.positions[symbol]["contracts"]
            * self.markets[symbol]["contractSize"]
            * self._get_convert_ratio(self.markets[symbol]["quote"], "USDT")
        )
        if self.positions[symbol]["side"] == "short":
            profit = -profit
        elif self.positions[symbol]["side"] == "long":
            profit = profit
        else:
            raise Exception(f"position side error : {self.positions[symbol]['side']}")
        return profit

    def get_unrealized_value(self):
        unrealized_value = 0
        for symbol in self.positions:
            unrealized_value += self.get_symbol_unrealized(symbol)
        return unrealized_value

    def get_all_balance_margin(self):
        spot_margin = 0
        for balance in self.balances:
            margin = self.get_balance_margin(balance, quote="USDT")
            spot_margin += margin
        return spot_margin

    def get_future_margin(self, symbol: str):
        if symbol not in self.positions:
            return 0
        margin = (
            self.positions[symbol]["contracts"]
            * self.markets[symbol]["contractSize"]
            * self._get_convert_ratio(self.markets[symbol]["base"], "USDT")
            / self.get_leverage(symbol)
        )
        return margin

    def get_all_future_margin(self):
        future_margin = 0
        for symbol in self.positions:
            margin = self.get_future_margin(symbol)
            future_margin += margin
        return future_margin

    def get_account_value(self):
        value = self.get_all_balance_value() + self.get_unrealized_value()
        return value

    def get_free_margin(self, unit: str) -> float:
        result = (
            self.get_all_balance_value()
            + self.get_unrealized_value()
            - self.get_all_balance_margin()
            - self.get_all_future_margin()
        ) / self._get_convert_ratio("USDT", unit)
        return result

    def get_unimrr(self):
        balance_value = self.get_all_balance_value()
        unrealized_value = self.get_unrealized_value()
        spot_margin = self.get_all_balance_margin()
        future_margin = self.get_all_future_margin()
        value = balance_value + unrealized_value
        unimrr = value / (
            (spot_margin + future_margin) + 1e-9
        )  # todo: 维持保证金算法错误
        return unimrr

    def _get_match_price(
        self, symbol: str, price: float, is_market: bool, is_over: bool
    ):
        cur_price = self._get_cur_price(symbol)
        if is_market:
            if is_over:
                cur_price *= 1 + self.slip_rate
            else:
                cur_price *= 1 - self.slip_rate
            return cur_price
        else:
            if is_over:
                if price >= self.get_data(symbol, "low", -1):
                    return price
                else:
                    return None
            else:
                if price <= self.get_data(symbol, "high", -1):
                    return price
                else:
                    return None

    def get_future_can_do(self, symbol: str, side: OrderSide = "buy"):
        free_margin = self.get_free_margin("USDT")
        cur_size = self.get_position_size(symbol)
        if cur_size == 0 and (
            (side == "buy" and cur_size > 0) or (side == "sell" and cur_size < 0)
        ):
            return (
                free_margin
                * self.get_leverage(symbol)
                / self._get_convert_ratio(self.markets[symbol]["base"], "USDT")
            )
        else:
            used_margin = self.get_future_margin(symbol)
            return (
                (free_margin + used_margin)
                * self.get_leverage(symbol)
                / self._get_convert_ratio(self.markets[symbol]["base"], "USDT")
            )
            # 未实现利润在free_margin计算时已经考虑到了

    def get_margin_can_do(self, symbol: str):
        """
        return base_buy, base_sell, quote_buy, quote_sell
        """
        free_margin = self.get_free_margin("USDT")
        base = self.markets[symbol]["base"]
        quote = self.markets[symbol]["quote"]
        self.get_balance_size(base)
        self.get_balance_size(quote)
        org = self.get_balance_size(quote) * self._get_convert_ratio(
            quote, "USDT"
        ) + self.get_balance_size(base) * self._get_convert_ratio(base, "USDT")
        rmv = (
            free_margin
            + self.get_balance_margin(base, quote="USDT")
            + self.get_balance_margin(quote, quote="USDT")
            - org
        )
        pa = self._get_cur_price(base)
        pb = self._get_cur_price(quote)
        ca = self.get_collateral(base)
        cb = self.get_collateral(quote)
        lva = self.get_balance_leverage(base)
        lvb = self.get_balance_leverage(quote)

        # 其他资产可抵用的保证金

        while True:
            # sa >= 0 sb >=0
            sa = (rmv + org * cb) / (pa * cb - pa * ca)
            sb = (org - sa * pa) / pb
            if sa >= 0 and sb >= 0:
                break

            # sa >= 0 sb < 0
            sa = (rmv + org / lvb + org * cb) / (pa / lvb - pa * ca + pa * cb)
            sb = (org - sa * pa) / pb
            if sa >= 0 and sb < 0:
                break

            # sa < 0 sb >= 0
            sb = (rmv + org / lva + org * ca) / (pb / lva - pb * cb + pb * ca)
            sa = (org - sb * pb) / pa

            if sa < 0 and sb >= 0:
                break

            # sa < 0 sb < 0
            sa = (rmv + org * cb + org / lvb) / (
                -pa / lva + pa / lvb + pa * cb - pa * ca
            )
            sb = (org - sa * pa) / pb
            if sa < 0 and sb < 0:
                break
            raise Exception(f"尝试获得{symbol}的可调整仓位失败")

    def force_clear_positions(self):
        for symbol in self.positions:
            self.adjust_to_target_future(symbol, 0, None, "GTC", "market")

    def force_clear_balances(self):
        for balance in self.balances:
            if balance != "USDT":
                self.adjust_to_target_spot(f"{balance}/USDT", 0, None, "GTC", "market")

    @abc.abstractmethod
    def get_cur_time(self):
        pass

    def order_future(
        self,
        symbol: str,
        size: float,
        price: float,
        timeInForce: str,
        market_type,
    ) -> bool:
        pass

    def adjust_to_target_future(
        self,
        symbol: str,
        size: float,
        price: float,
        timeInForce: str,
        market_type,
    ):
        old_value = self.get_position_size(symbol)
        adjust = size - old_value
        self.order_future(symbol, adjust, price, timeInForce, market_type)

    def order_spot(
        self,
        symbol: str,
        size: float,
        price: float,
        timeInForce: str,
        market_type: str,
    ):
        pass

    def adjust_to_target_spot(
        self,
        symbol: str,
        value: float,
        price: float,
        timeInForce: str,
        market_type: str,
    ):
        quote = self.markets[symbol]["base"]
        old_value = self.get_balance_size(self.balances[quote])
        adjust = value - old_value
        if adjust != 0:
            return self.order_spot(
                symbol,
                adjust,
                price,
                timeInForce,
                market_type,
            )
        return None
